{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1635a6fd",
   "metadata": {},
   "source": [
    "[download this notebook here](https://github.com/HumanCompatibleAI/imitation/blob/master/docs/tutorials/5a_train_preference_comparisons_with_cnn.ipynb)\n",
    "\n",
    "# Learning a Reward Function using Preference Comparisons on Atari\n",
    "\n",
    "In this case, we will use a convolutional neural network for our policy and reward model. We will also shape the learned reward model with the policy's learned value function, since these shaped rewards will be more informative for training - incentivizing agents to move to high-value states. In the interests of execution time, we will only do a little bit of training - much less than in the previous preference comparison notebook. To run this notebook, be sure to install the `atari` extras, for example by running `pip install imitation[atari]`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08bdff31",
   "metadata": {},
   "source": [
    "First, we will set up the environment, reward network, et cetera."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93187e19",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch as th\n",
    "import gymnasium as gym\n",
    "from gymnasium.wrappers import TimeLimit\n",
    "import numpy as np\n",
    "\n",
    "from seals.util import AutoResetWrapper\n",
    "\n",
    "from stable_baselines3 import PPO\n",
    "from stable_baselines3.common.atari_wrappers import AtariWrapper\n",
    "from stable_baselines3.common.env_util import make_vec_env\n",
    "from stable_baselines3.common.vec_env import VecFrameStack\n",
    "from stable_baselines3.ppo import CnnPolicy\n",
    "\n",
    "from imitation.algorithms import preference_comparisons\n",
    "from imitation.data.wrappers import RolloutInfoWrapper\n",
    "from imitation.policies.base import NormalizeFeaturesExtractor\n",
    "from imitation.rewards.reward_nets import CnnRewardNet\n",
    "\n",
    "\n",
    "device = th.device(\"cuda\" if th.cuda.is_available() else \"cpu\")\n",
    "\n",
    "rng = np.random.default_rng()\n",
    "\n",
    "\n",
    "# Here we ensure that our environment has constant-length episodes by resetting\n",
    "# it when done, and running until 100 timesteps have elapsed.\n",
    "# For real training, you will want a much longer time limit.\n",
    "def constant_length_asteroids(num_steps):\n",
    "    atari_env = gym.make(\"AsteroidsNoFrameskip-v4\")\n",
    "    preprocessed_env = AtariWrapper(atari_env)\n",
    "    endless_env = AutoResetWrapper(preprocessed_env)\n",
    "    limited_env = TimeLimit(endless_env, max_episode_steps=num_steps)\n",
    "    return RolloutInfoWrapper(limited_env)\n",
    "\n",
    "\n",
    "# For real training, you will want a vectorized environment with 8 environments in parallel.\n",
    "# This can be done by passing in n_envs=8 as an argument to make_vec_env.\n",
    "# The seed needs to be set to 1 for reproducibility and also to avoid win32\n",
    "# np.random.randint high bound error.\n",
    "venv = make_vec_env(constant_length_asteroids, env_kwargs={\"num_steps\": 100}, seed=1)\n",
    "venv = VecFrameStack(venv, n_stack=4)\n",
    "\n",
    "reward_net = CnnRewardNet(\n",
    "    venv.observation_space,\n",
    "    venv.action_space,\n",
    ").to(device)\n",
    "\n",
    "fragmenter = preference_comparisons.RandomFragmenter(warning_threshold=0, rng=rng)\n",
    "gatherer = preference_comparisons.SyntheticGatherer(rng=rng)\n",
    "preference_model = preference_comparisons.PreferenceModel(reward_net)\n",
    "reward_trainer = preference_comparisons.BasicRewardTrainer(\n",
    "    preference_model=preference_model,\n",
    "    loss=preference_comparisons.CrossEntropyRewardLoss(),\n",
    "    epochs=3,\n",
    "    rng=rng,\n",
    ")\n",
    "\n",
    "agent = PPO(\n",
    "    policy=CnnPolicy,\n",
    "    env=venv,\n",
    "    seed=0,\n",
    "    n_steps=16,  # To train on atari well, set this to 128\n",
    "    batch_size=16,  # To train on atari well, set this to 256\n",
    "    ent_coef=0.01,\n",
    "    learning_rate=0.00025,\n",
    "    n_epochs=4,\n",
    ")\n",
    "\n",
    "trajectory_generator = preference_comparisons.AgentTrainer(\n",
    "    algorithm=agent,\n",
    "    reward_fn=reward_net,\n",
    "    venv=venv,\n",
    "    exploration_frac=0.0,\n",
    "    rng=rng,\n",
    ")\n",
    "\n",
    "pref_comparisons = preference_comparisons.PreferenceComparisons(\n",
    "    trajectory_generator,\n",
    "    reward_net,\n",
    "    num_iterations=2,\n",
    "    fragmenter=fragmenter,\n",
    "    preference_gatherer=gatherer,\n",
    "    reward_trainer=reward_trainer,\n",
    "    fragment_length=10,\n",
    "    transition_oversampling=1,\n",
    "    initial_comparison_frac=0.1,\n",
    "    allow_variable_horizon=False,\n",
    "    initial_epoch_multiplier=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ceadb25",
   "metadata": {},
   "source": [
    "We are now ready to train the reward model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c2c4d3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "pref_comparisons.train(\n",
    "    total_timesteps=16,\n",
    "    total_comparisons=15,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5de1f5f",
   "metadata": {},
   "source": [
    "We can now wrap the environment with the learned reward model, shaped by the policy's learned value function. Note that if we were training this for real, we would want to normalize the output of the reward net as well as the value function, to ensure their values are on the same scale. To do this, use the `NormalizedRewardNet` class from `src/imitation/rewards/reward_nets.py` on `reward_net`, and modify the potential to add a `RunningNorm` module from `src/imitation/util/networks.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbcee105",
   "metadata": {},
   "outputs": [],
   "source": [
    "from imitation.rewards.reward_nets import ShapedRewardNet, cnn_transpose\n",
    "from imitation.rewards.reward_wrapper import RewardVecEnvWrapper\n",
    "\n",
    "\n",
    "def value_potential(state):\n",
    "    state_ = cnn_transpose(state)\n",
    "    return agent.policy.predict_values(state_)\n",
    "\n",
    "\n",
    "shaped_reward_net = ShapedRewardNet(\n",
    "    base=reward_net,\n",
    "    potential=value_potential,\n",
    "    discount_factor=0.99,\n",
    ")\n",
    "\n",
    "# GOTCHA: When using the NormalizedRewardNet wrapper, you should deactivate updating\n",
    "# during evaluation by passing update_stats=False to the predict_processed method.\n",
    "learned_reward_venv = RewardVecEnvWrapper(venv, shaped_reward_net.predict_processed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41efea37",
   "metadata": {},
   "source": [
    "Next, we train an agent that sees only the shaped, learned reward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63cc3121",
   "metadata": {},
   "outputs": [],
   "source": [
    "learner = PPO(\n",
    "    policy=CnnPolicy,\n",
    "    env=learned_reward_venv,\n",
    "    seed=0,\n",
    "    batch_size=64,\n",
    "    ent_coef=0.0,\n",
    "    learning_rate=0.0003,\n",
    "    n_epochs=10,\n",
    "    n_steps=64,\n",
    ")\n",
    "learner.learn(1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16331df8",
   "metadata": {},
   "source": [
    "We now evaluate the learner using the original reward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2c4828b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from stable_baselines3.common.evaluation import evaluate_policy\n",
    "\n",
    "reward, _ = evaluate_policy(learner.policy, venv, 10)\n",
    "print(reward)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hZ9um7u1V0h1",
   "metadata": {},
   "source": [
    "## Generating rollouts\n",
    "When generating rollouts in image environments, be sure to use the agent's get_env() function rather than using the original environment.\n",
    "\n",
    "The learner re-arranges the observations space to put the channel environment in the first dimension, and get_env() will correctly provide a wrapped environment doing this.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bpK6ih-3WjqG",
   "metadata": {},
   "outputs": [],
   "source": [
    "from imitation.data import rollout\n",
    "\n",
    "rollouts = rollout.rollout(\n",
    "    learner,\n",
    "    # Note that passing venv instead of agent.get_env()\n",
    "    # here would fail.\n",
    "    learner.get_env(),\n",
    "    rollout.make_sample_until(min_timesteps=None, min_episodes=3),\n",
    "    rng=rng,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
